# -*- coding: utf-8 -*-
import torch
from torch import nn
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader
from dataset import CustomDataset
from model import CustomNet
import os
from torchvision import transforms


def train_loop(epochs, dataloader, model, loss_fn, optimizer, device):
    """训练流程"""
    model.train()
    best_acc = 0.0

    for epoch in range(epochs):
        running_loss = 0.0
        correct = 0
        total = 0

        for batch_idx, data in enumerate(dataloader):
            inputs = data['image'].to(device)
            labels = data['label'].to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_fn(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            if batch_idx % 10 == 0:
                print(f'Epoch [{epoch + 1}/{epochs}], Batch [{batch_idx + 1}/{len(dataloader)}], '
                      f'Loss: {running_loss / 10:.4f}')
                running_loss = 0.0

        epoch_acc = 100 * correct / total
        print(f'Epoch {epoch + 1} Accuracy: {epoch_acc:.2f}%')

        if epoch_acc > best_acc:
            best_acc = epoch_acc
            torch.save(model.state_dict(), './models/best_model.pth')
            print(f'Best model saved with accuracy: {best_acc:.2f}%')


def main():
    # 超参数设置
    BATCH_SIZE = 32
    LEARNING_RATE = 0.001
    EPOCHS = 50

    # 设备配置
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    try:
        # 数据增强
        train_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(10),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])

        # 修正路径 - 使用相对路径或正确绝对路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        train_txt_path = os.path.join(current_dir, 'images', 'train.txt')
        train_img_dir = os.path.join(current_dir, 'images', 'train')

        print(f"Loading data from:\n{train_txt_path}\n{train_img_dir}")

        # 数据加载
        train_dataset = CustomDataset(
            annotations_file=train_txt_path,
            img_dir=train_img_dir,
            transform=lambda: train_transform
        )
        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

        # 自动获取类别数
        num_classes = len(set([int(label) for _, label in train_dataset.labels]))
        print(f"Detected {num_classes} classes in dataset")

        # 模型初始化
        model = CustomNet(num_classes=num_classes).to(device)

        # 损失函数和优化器
        criterion = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

        # 创建模型目录
        os.makedirs('./models', exist_ok=True)

        # 开始训练
        print("Starting training...")
        train_loop(EPOCHS, train_loader, model, criterion, optimizer, device)

    except Exception as e:
        print(f"Error during training: {str(e)}")
        print("请检查以下内容：")
        print("1. 确保 images/train.txt 和 images/train/ 目录存在")
        print("2. train.txt 文件内容格式为：'图片路径 标签'（每行一个样本）")
        print("3. 图片实际存在于指定路径")


if __name__ == "__main__":
    main()